
#include <windows.h>

#include <Objbase.h> // for CoInitializeEx() / CoUninitialize() / CoCreateInstance()

#include <mfapi.h> // for MFStartup() / MFShutdown() / MFCreateAttributes() / MFTEnum()
#include <Mfidl.h> // for MFCreatePresentationClock() / MFCreateAudioRenderer()
#include <mfreadwrite.h> // for MFCreateSourceReaderFromURL()
#include <mmdeviceapi.h> // for IMMDeviceEnumerator / IMMDeviceCollection 

#include <mftransform.h>
#include <mferror.h>

#include <list>

#include <stdio.h>

#include "audio_util.hpp"

IMFMediaSink* GetAudioSinkDevice(void)
{
    HRESULT hr;

    IMFMediaSink *audio_sink = NULL;             // Streaming audio renderer (SAR)

    IMMDeviceEnumerator *audio_render_device_enum = NULL;      // Audio device enumerator.
    IMMDeviceCollection *audio_devices = NULL;   // Audio device collection.
    LPWSTR wstrID = NULL;                   // Device ID.
    IMFAttributes *pAttributes = NULL;      // Attribute store.

    do
    {
        // Create the device enumerator.
        hr = CoCreateInstance(__uuidof(MMDeviceEnumerator), NULL, CLSCTX_ALL, __uuidof(IMMDeviceEnumerator), (void**)&audio_render_device_enum);
        if (hr != S_OK)
        {
            break;
        }

         // Enumerate the rendering devices.
        hr = audio_render_device_enum->EnumAudioEndpoints(eRender, DEVICE_STATE_ACTIVE, &audio_devices);
        if (hr != S_OK)
        {
            break;
        }

        // Get ID of the first device in the list.
        IMMDevice *device = NULL;              // An audio device.
        hr = audio_devices->Item(0, &device);
        if (hr != S_OK)
        {
            break;
        }

        hr = device->GetId(&wstrID);
        if (hr != S_OK)
        {
            break;
        }

        // Create an attribute store and set the device ID attribute.
        hr = MFCreateAttributes(&pAttributes, 2);
        if (hr != S_OK)
        {
            break;
        }

        hr = pAttributes->SetString(MF_AUDIO_RENDERER_ATTRIBUTE_ENDPOINT_ID, wstrID);
        // Create the audio renderer.
        if (hr != S_OK)
        {
            break;
        }

        hr = MFCreateAudioRenderer(pAttributes, &audio_sink); 
    } while(0);

	SafeRelease(audio_render_device_enum);
	SafeRelease(audio_devices);
	SafeRelease(pAttributes);
    CoTaskMemFree(wstrID);

    return audio_sink;
}

void SetupAudioSink
(
    IMFMediaSink *sink, 
    IMFMediaType *input_media_type,
    IMFPresentationClock *present_clock,
    IMFStreamSink **stream_sink_out,
    DWORD *sink_stream_index_out,
    DWORD *sink_spec_out

)
{
    HRESULT hr;

    hr = sink->SetPresentationClock(present_clock);

    DWORD sink_spec = 0;
    sink->GetCharacteristics(&sink_spec);

    DWORD sink_stream_index = (DWORD)-1;
    IMFStreamSink *stream_sink = NULL;
    if (sink_spec & MEDIASINK_FIXED_STREAMS)
    {
        sink_stream_index = 0;
        hr = sink->GetStreamSinkByIndex(0, &stream_sink);
    }
    else
    {
        sink_stream_index = 0;
        hr = sink->AddStreamSink(0, NULL, &stream_sink);
    }

    IMFMediaTypeHandler *media_type_handler = NULL;
    hr = stream_sink->GetMediaTypeHandler(&media_type_handler);
#if 0
    {
        DWORD idx = 0;
        do
        {
            IMFMediaType *type = NULL;
            hr = media_type_handler->GetMediaTypeByIndex(idx, &type);
            if (hr != S_OK) { break; }

            DumpAudioMediaType(type);
            type->Release();

            idx++;
        } while (1);
    }
#endif

    hr = media_type_handler->SetCurrentMediaType(input_media_type);
    if (hr != S_OK)
    {
        IMFMediaType *media_type_audio_play = NULL;
        hr = media_type_handler->IsMediaTypeSupported(input_media_type, &media_type_audio_play);
        if (hr == S_OK && media_type_audio_play != NULL)
        {
            media_type_handler->SetCurrentMediaType(media_type_audio_play);
        }
        SafeRelease(media_type_audio_play);
    }
    SafeRelease(media_type_handler);

    IMFMediaSinkPreroll *media_sink_preroll = NULL;
    if (sink_spec & MEDIASINK_CAN_PREROLL)
    {
        hr = sink->QueryInterface(IID_PPV_ARGS(&media_sink_preroll));
    }
    if (media_sink_preroll != NULL)
    {
        media_sink_preroll->NotifyPreroll(1000000);
        media_sink_preroll->Release();
    }
    else
    {
        present_clock->Start(1000000);
    }

    *stream_sink_out = stream_sink;
    *sink_stream_index_out = sink_stream_index;
    *sink_spec_out = sink_spec;

    return;
}

int CreateSourceReaderFromFile
(
    const wchar_t *file, 
    IMFSourceReader **reader, 
    DWORD *audio_stream_index_out, IMFMediaType **audio_media_type_out,
    DWORD *video_stream_index_out, IMFMediaType **video_media_type_out
)
{
	HRESULT hr;

    IMFSourceReader *media_reader = NULL;

  #if 1
    hr = MFCreateSourceReaderFromURL(file, NULL, &media_reader);
  #else
    {
        IMFMediaSource *media_source = NULL;

        IMFSourceResolver *resolver = NULL;
        hr = MFCreateSourceResolver(&resolver);

        MF_OBJECT_TYPE obj_type = MF_OBJECT_INVALID;
        IUnknown *out_obj = NULL;
        hr = resolver->CreateObjectFromURL(file, MF_RESOLUTION_MEDIASOURCE, NULL, &obj_type, &out_obj);

        hr = out_obj->QueryInterface(IID_PPV_ARGS(&media_source));

        hr = MFCreateSourceReaderFromMediaSource(media_source, NULL, &media_reader);

        SafeRelease(media_source);
        SafeRelease(resolver);
        SafeRelease(out_obj);
    }
  #endif

    DWORD audio_stream_index = (DWORD)-1;
    IMFMediaType *media_type_audio = NULL;

    DWORD video_stream_index = (DWORD)-1;
    IMFMediaType *media_type_video = NULL;

    /* enumerate all media stream type */
    DWORD stream_index = 0;
    do
    {
        IMFMediaType *media_type = NULL;
        hr = media_reader->GetCurrentMediaType(stream_index, &media_type);
        if (hr == S_OK)
        {
            GUID major_type;
            media_type->GetMajorType(&major_type);
            if (audio_stream_index == (DWORD)-1 && major_type == MFMediaType_Audio)
            {
                audio_stream_index = stream_index;
                media_type_audio = media_type;
            }
            else if (video_stream_index == (DWORD)-1 && major_type == MFMediaType_Video)
            {
                video_stream_index = stream_index;
                media_type_video = media_type;
            }
            else
            {
                media_type->Release();
            }

            stream_index++;
        }
    } while (hr == S_OK);

    hr = media_reader->SetCurrentMediaType(audio_stream_index, NULL, media_type_audio);
    hr = media_reader->SetCurrentMediaType(video_stream_index, NULL, media_type_video);

    *reader = media_reader;

    *audio_stream_index_out = audio_stream_index;
    *audio_media_type_out = media_type_audio;

    *video_stream_index_out = video_stream_index;
    *video_media_type_out = media_type_video;

	return 0;
}

IMFSample *CreateDecodeOutputSample(DWORD buf_size)
{
    if (buf_size == 0) { return NULL; }

    IMFSample *sample = NULL;
    MFCreateSample(&sample);

    IMFMediaBuffer *media_buf = NULL;
    MFCreateMemoryBuffer(buf_size, &media_buf);
    sample->AddBuffer(media_buf);
    media_buf->Release();

    return sample;
}

void DestroyDecodeOutputSample(IMFSample *sample)
{
    DWORD buf_count = 0;
    sample->GetBufferCount(&buf_count);
    for (DWORD buf_idx = 0; buf_idx < buf_count; ++buf_idx)
    {
        IMFMediaBuffer *buf = NULL;
        HRESULT hr = sample->GetBufferByIndex(buf_idx, &buf);
        if (buf != NULL)
        {
            buf->Release();
        }
        //sample->RemoveBufferByIndex(buf_idx);
    }
    sample->Release();

    return;
}

int CreateAudioDecoder(IMFMediaType *input_media_type, IMFTransform **audio_decoder, IMFMediaType **output_media_type, DWORD *input_id, DWORD *output_id, DWORD *output_buf_size)
{
    HRESULT hr;
    IMFTransform *decoder = NULL;

    GUID sub_type;
    input_media_type->GetGUID(MF_MT_SUBTYPE, &sub_type);

    MFT_REGISTER_TYPE_INFO in_type_info = { MFMediaType_Audio, sub_type };
    MFT_REGISTER_TYPE_INFO out_type_info = { MFMediaType_Audio, MFAudioFormat_Float };

    CLSID *clsids = NULL;
    UINT32 clsids_count = 0;
    hr = MFTEnum(MFT_CATEGORY_AUDIO_DECODER, 0, &in_type_info, &out_type_info, NULL, &clsids, &clsids_count);
    if (hr != S_OK || clsids_count == 0)
    {
        wprintf(L"failed to enum auido decoder, hr: %d\n", hr);
        return -1;
    }

    hr = CoCreateInstance(clsids[0], NULL, CLSCTX_INPROC_SERVER, IID_PPV_ARGS(&decoder));

    CoTaskMemFree(clsids);

    DWORD input_stream_count = 0;
    DWORD outut_stream_count = 0;
    hr = decoder->GetStreamCount(&input_stream_count, &outut_stream_count);

    DWORD input_stream_id = (DWORD)-1;
    DWORD output_stream_id = (DWORD)-1;
    if (input_stream_count < 1 || outut_stream_count < 1)
    {
        decoder->Release();
        decoder = NULL;
    }
    else
    {
        DWORD in_out_stream_count = input_stream_count + outut_stream_count;
        DWORD *stream_ids = new DWORD[in_out_stream_count];
        DWORD *input_stream_ids = stream_ids;
        DWORD *output_stream_ids = stream_ids + input_stream_count;
        hr = decoder->GetStreamIDs(input_stream_count, input_stream_ids, outut_stream_count, output_stream_ids);
        if (hr == S_OK)
        {
            input_stream_id = input_stream_ids[0];
            output_stream_id = output_stream_ids[0];
        }
        else if (hr == E_NOTIMPL)
        {
            input_stream_id = 0;
            output_stream_id = 0;
        }

        delete []stream_ids;
    }

    if (input_stream_id == (DWORD)-1 || output_stream_id == (DWORD)-1)
    {
        wprintf(L"failed to get in/out stream ids, hr: %d\n", hr);
        decoder->Release();
        decoder = NULL;

        return -2;
    }

    hr = decoder->SetInputType(input_stream_id, input_media_type, 0);

    DWORD out_media_type_index = 0;
    do
    { 
        IMFMediaType *out_media_type = NULL;
        hr = decoder->GetOutputAvailableType(output_stream_id, out_media_type_index, &out_media_type);
        if (hr != S_OK) { break; }
        //DumpAudioMediaType(out_media_type);

        GUID sub_type;
        out_media_type->GetGUID(MF_MT_SUBTYPE, &sub_type);
        if (sub_type == MFAudioFormat_Float)
        {
            *output_media_type = out_media_type;
            break;
        }
        out_media_type->Release();

        out_media_type_index++;
    } while (1);
    hr = decoder->SetOutputType(output_stream_id, *output_media_type, 0);

    MFT_INPUT_STREAM_INFO input_stream_info;
    hr = decoder->GetInputStreamInfo(input_stream_id, &input_stream_info);

    MFT_OUTPUT_STREAM_INFO output_stream_info;
    hr = decoder->GetOutputStreamInfo(output_stream_id, &output_stream_info);

    *audio_decoder = decoder;
    *input_id = input_stream_id;
    *output_id = output_stream_id;
    if (output_stream_info.dwFlags & (MFT_OUTPUT_STREAM_PROVIDES_SAMPLES | MFT_OUTPUT_STREAM_CAN_PROVIDE_SAMPLES))
    {
        *output_buf_size = 0;
    }
    else
    {
        *output_buf_size = output_stream_info.cbSize;
    }

    return 0;
}

void AudioPlay(const wchar_t *file)
{
    HRESULT hr;

    IMFSourceReader *reader = NULL;
    DWORD audio_stream_index = (DWORD)-1;
    IMFMediaType *media_type_audio = NULL;
    DWORD video_stream_index = (DWORD)-1;
    IMFMediaType *media_type_video = NULL;

    CreateSourceReaderFromFile(file, &reader, &audio_stream_index, &media_type_audio, &video_stream_index, &media_type_video);

    IMFTransform *audio_decoder = NULL;
    IMFMediaType *decoded_audio_media_type = NULL;
    DWORD decode_input_id, decode_output_id;
    DWORD decode_output_buf_size = 0;
    CreateAudioDecoder(media_type_audio, &audio_decoder, &decoded_audio_media_type, &decode_input_id, &decode_output_id, &decode_output_buf_size);

    // ==================================================

    IMFPresentationClock *present_clock = NULL;
    hr = MFCreatePresentationClock(&present_clock);
    IMFPresentationTimeSource *clock_source = NULL;
    MFCreateSystemTimeSource(&clock_source);
    hr = present_clock->SetTimeSource(clock_source);
    SafeRelease(clock_source);

    IMFMediaSink *sink = GetAudioSinkDevice();
    IMFStreamSink *stream_sink = NULL;
    DWORD sink_stream_index = (DWORD)-1;
    DWORD sink_spec = 0;
    SetupAudioSink(sink, decoded_audio_media_type, present_clock, &stream_sink, &sink_stream_index, &sink_spec);

    // ==================================================

    hr = audio_decoder->ProcessMessage(MFT_MESSAGE_NOTIFY_BEGIN_STREAMING, 0);

    bool source_stream_eof = false;
    bool decode_stream_eof = false;
    std::list<IMFSample*> pending_samples;
    int pending_samples_count = 0;
    for (;;)
    {
        if (!source_stream_eof && pending_samples_count < 2)
        {
            DWORD streamIndex, flags = 0;
            LONGLONG llTimeStamp;
            IMFSample *sample = NULL;
            hr = reader->ReadSample(audio_stream_index, 0, &streamIndex, &flags, &llTimeStamp, &sample);
            if (sample != NULL)
            {
                hr = audio_decoder->ProcessInput(decode_input_id, sample, 0);
                sample->Release();
            }

            if (flags == MF_SOURCE_READERF_ENDOFSTREAM)
            {
                wprintf(L"All source data has been processed\n");
                source_stream_eof = true;

                audio_decoder->ProcessMessage(MFT_MESSAGE_NOTIFY_END_OF_STREAM, 0);
                audio_decoder->ProcessMessage(MFT_MESSAGE_COMMAND_DRAIN, 0);
            }
        }

        if (!decode_stream_eof && pending_samples_count < 2)
        {
            MFT_OUTPUT_DATA_BUFFER decode_output = {decode_output_id, NULL, 0, NULL};
            decode_output.pSample = CreateDecodeOutputSample(decode_output_buf_size);
            DWORD output_flags = 0;
            hr = audio_decoder->ProcessOutput(0, 1, &decode_output, &output_flags);
            if (hr == S_OK)
            {
                pending_samples.push_back(decode_output.pSample);
                pending_samples_count++;
            }
            else 
            {
                DestroyDecodeOutputSample(decode_output.pSample);
                if (hr == MF_E_TRANSFORM_NEED_MORE_INPUT)
                {
                    if (source_stream_eof)
                    {
                        wprintf(L"All decoder data has been processed\n");
                        decode_stream_eof = true;
                    }
                }
            }
        }

        if (decode_stream_eof && pending_samples_count == 0)
        {
            break;
        }

        IMFMediaEvent *event = NULL;
        hr = stream_sink->GetEvent(MF_EVENT_FLAG_NO_WAIT, &event);
        if (hr != S_OK) { continue; }

        MediaEventType event_type;
        event->GetType(&event_type);
        switch (event_type)
        {
            case MEStreamSinkRequestSample:
            {
                if (!pending_samples.empty())
                {
                    IMFSample *sample = pending_samples.front();
                    pending_samples.pop_front();
                    pending_samples_count--;

                    stream_sink->ProcessSample(sample);
                    DestroyDecodeOutputSample(sample);
                }

                break;
            }
            case MEStreamSinkPrerolled:
            {
                present_clock->Start(1000000);
                break;
            }
            default: break;
        }

        event->Release();
    }

    present_clock->Stop();
    SafeRelease(present_clock);

    if ((sink_spec & MEDIASINK_FIXED_STREAMS) == 0) { sink->RemoveStreamSink(sink_stream_index); }
    SafeRelease(stream_sink);
    SafeRelease(sink);

    SafeRelease(decoded_audio_media_type);
    SafeRelease(audio_decoder);

    SafeRelease(reader);
    SafeRelease(media_type_audio);
    SafeRelease(media_type_video);

    return;
}

/************************************************************/

void audio_play_routine(int argc, wchar_t *argv[])
{
    CoInitializeEx(0, COINIT_MULTITHREADED);
    MFStartup(MF_VERSION);

	AudioPlay(argv[1]);

    MFShutdown();
    CoUninitialize();
    
    return;
}

int wmain(int argc, wchar_t *argv[])
{
    audio_play_routine(argc, argv);

    return 0;
}

